import numpy as np
import time
import matplotlib.pyplot as plt

# SOR iteration
def solve_SOR(A, b, x0, K, e, w):
    n = len(A)
    x0 = x0.flatten()
    k = 0  # Iteration count
    err = np.inf  # Initialize error as infinity

    while k < K and err > e:
        k += 1
        x1 = x0.copy()
        
        for i in range(n):
            sum1 = np.dot(A[i, :i], x0[:i])    # Calculate upper triangular part
            sum2 = np.dot(A[i, i+1:], x1[i+1:])  # Calculate lower triangular part
            x0[i] = (1 - w) * x1[i] + (w / A[i][i]) * (b[i] - sum1 - sum2)

        # Calculate error (maximum norm between k+1 and k)
        err = np.max(np.abs(x1 - x0))

    return k, x0.reshape(-1, 1)

if __name__ == "__main__":
    # Load parameters from input file
    with open('input.txt', 'r') as f:
        lines = f.readlines()

    # Remove comments and empty lines
    lines = [line for line in lines if line.strip() and not line.strip().startswith('#')]

    # Read matrix A
    A = np.array([list(map(float, line.split())) for line in lines[0:9]])
    # Read vector b
    b = np.array([float(lines[i].strip()) for i in range(9, 18)]).reshape(-1, 1)
    # Read initial guess x0
    x0 = np.array([float(lines[i].strip()) for i in range(18, 27)]).reshape(-1, 1)
    # Read maximum iterations K
    K = int(lines[27].strip())
    # Read error tolerance epsilon
    e = float(lines[28].strip())
    # Read relaxation factor w
    w = float(lines[29].strip())

    # Timing
    start_time = time.time()
    
    # SOR iteration
    k, x = solve_SOR(A, b, x0, K, e, w)
    end_time = time.time()
    print("SOR Iteration")
    print("Number of iterations:", k)
    print("Approximate solution:\n", x)
    print("Elapsed time: {:.4f} seconds".format(end_time - start_time))

    # Gauss-Seidel iteration (using SOR with w=1)
    x0 = np.zeros((9, 1))
    k_gs, x_gs = solve_SOR(A, b, x0, K, e, 1)
    print("Gauss-Seidel Iteration")
    print("Number of iterations:", k_gs)
    print("Approximate solution:\n", x_gs)
